[Common][PyTorch] EP dispatch with unfused MXFP8 quantization by phu0ngng · Pull Request #3270 · NVIDIA/TransformerEngine

phu0ngng · 2026-07-28T20:33:22Z

Description

This PR adds MXFP8 support to the dispatch op of the NCCL EP path. The dispatch op is used in two places, and MXFP8 applies to both:

Dispatch forward bfloat16 tokens are quantized to MXFP8 internally and dispatched to the target experts; recv is returned as a per-expert GroupedTensor.
Combine backward the result-grad is scattered back to expert positions through the same (reverse) dispatch op, quantized to MXFP8, returning the expert-output grad as a per-expert GroupedTensor.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

PyTorch frontend (`transformer_engine/pytorch/ep.py`, `distributed.py`, `csrc/extensions/ep.cpp`)**

The dispatch op quantizes bfloat16 tokens to MXFP8 internally when the buffer's dispatch_quant_recipe is set (MXFP8BlockScaling only for now); dispatch-forward recv is returned as a per-expert GroupedTensor. A pre-quantized input is rejected.
Combine backward reuses the dispatch op to scatter the result-grad: it quantizes the grad to MXFP8 and returns the expert-output grad as a per-expert GroupedTensor. Combine forward is unchanged (high-precision).
Recv data and block scales share a single caller-supplied (optionally symm-mem-backed) buffer, sliced into data-then-scale regions; the same convention is used for the combine backward grad buffer.

Common backend (`common/ep/ep_backend.cpp`, `include/.../ep.h`, `comm_window.h`)**

Backend and public headers extended to carry block-scale buffers/windows through the dispatch primitive.

NCCL EP submodule**

Bumped 3rdparty/nccl-extensions to the revision providing block-scaled dispatch.

Tests (`tests/cpp_distributed/test_ep.cu`, `tests/pytorch/distributed/run_ep.py`, `run_test_ep.sh`)**

Added C++ distributed coverage for the MXFP8 dispatch path.
Added PyTorch MXFP8 test passes for dispatch forward (normal, zero-copy, eager IO modes) and combine backward, gated behind a dedicated NVTE_EP_MXFP8_PASS run since the grouped path pins the per-expert alignment process-wide.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

greptile-apps · 2026-07-28T20:36:00Z

Greptile Summary

Adds unfused MXFP8 quantization support to NCCL expert-parallel dispatch and combine backward.

Routes MXFP8 data and block scales through common and PyTorch dispatch interfaces.
Returns per-expert quantized outputs as grouped tensors and supports caller-provided or symmetric-memory-backed buffers.
Extends distributed tests for normal, zero-copy, eager, and combine-backward paths.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains, and the previously reported ignore-file deletion has been fully restored in the current tree.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ep.py	Adds MXFP8 recipe handling, compact scale-buffer allocation, grouped output construction, and quantized combine-backward dispatch.
transformer_engine/pytorch/csrc/extensions/ep.cpp	Extends PyTorch bindings to validate, describe, and route MXFP8 data and scale-inverse tensors.
transformer_engine/common/ep/ep_backend.cpp	Adds NCCL descriptors and dispatch configuration for block-scaled token data and scales.
transformer_engine/pytorch/distributed.py	Adds explicit symmetric-memory pool lifecycle and cache cleanup needed by EP buffers.
tests/pytorch/distributed/run_ep.py	Adds MXFP8 dispatch and combine-backward coverage across fixed, eager, caller-buffer, and zero-copy modes.

Sequence Diagram

sequenceDiagram
  participant Caller
  participant PyEP as PyTorch EP
  participant Quant as MXFP8 Quantizer
  participant Backend as NCCL EP Backend
  participant Expert
  Caller->>PyEP: ep_dispatch(BF16 tokens)
  PyEP->>Quant: quantize data and block scales
  Quant-->>PyEP: E4M3 data + E8M0 scales
  PyEP->>Backend: dispatch data and scales
  Backend-->>Expert: per-expert GroupedTensor
  Expert->>PyEP: ep_combine(expert output)
  PyEP-->>Caller: high-precision combined result
  Caller->>PyEP: result gradient
  PyEP->>Quant: quantize gradient
  PyEP->>Backend: reverse dispatch data and scales
  Backend-->>Expert: grouped expert-output gradient

_{Reviews (10): Last reviewed commit: "Merge branch 'main' into ep_mxfp8" | Re-trigger Greptile}

phu0ngng · 2026-07-30T17:42:07Z

/te-ci L1 pytorch

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…der CUDA graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…CUDA-graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng · 2026-08-06T00:37:09Z

/te-ci L1

phu0ngng · 2026-08-06T20:06:15Z

/te-ci L1

phu0ngng · 2026-08-07T15:22:30Z

/te-ci L1

YangFei1990

Need to further align on API contracts before merging

phu0ngng · 2026-08-10T23:39:39Z

/te-ci L1

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng · 2026-08-11T16:11:25Z

/te-ci L1

phu0ngng requested review from ksivaman and ptrendx as code owners July 28, 2026 20:33

phu0ngng requested review from YangFei1990 and vthumbe1503 July 28, 2026 20:35

phu0ngng requested a review from zhongbozhu July 28, 2026 23:33

phu0ngng commented Jul 28, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ep.py Outdated

YangFei1990 requested changes Aug 2, 2026

View reviewed changes

phu0ngng force-pushed the ep_mxfp8 branch from 041b6e6 to 778c2c8 Compare August 4, 2026 16:29

YangFei1990 reviewed Aug 5, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ep.py Outdated

phu0ngng added 11 commits August 5, 2026 17:18

EP Dispatch with MXFP8

81ba2c5

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

restore gitignore

e4a572d

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

EP Dispatch with MXFP8

df4aaff

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Separate EP dispatch-forward and combine-backward quant recipes

01c7e1c

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Simplify EP quantize path and guard zero-copy pool alloc un…

a6f54a6

…der CUDA graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Add eager-mode coverage for MXFP8 EP combine backward

196419b

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Supply persistent symm-mem recv buffers for zero-copy 1F1B …

8934d2b

…CUDA-graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

shorten quant enum

ea10ceb

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Enable MXFP8 EP combine backward under zero-copy

b27bf5c

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Release symm-mem pool in ep_finalize

40f1bc8

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Drop unused device param from release_symm_mem_pool

1c30b1c

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng force-pushed the ep_mxfp8 branch from 778c2c8 to 1c30b1c Compare August 6, 2026 00:18

Merge branch 'main' into ep_mxfp8

dc05a03

Merge branch 'main' into ep_mxfp8

0ab4057

YangFei1990 approved these changes Aug 7, 2026

View reviewed changes

YangFei1990 self-requested a review August 10, 2026 04:13

YangFei1990 requested changes Aug 10, 2026

View reviewed changes

Merge branch 'main' into ep_mxfp8

a59f6f0

YangFei1990 previously approved these changes Aug 10, 2026

View reviewed changes

[PyTorch] Skip MXFP8 EP tests on pre-Blackwell devices

7cd92e0

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng dismissed YangFei1990’s stale review via 7cd92e0 August 11, 2026 16:10

ptrendx assigned timmoon10 Aug 11, 2026

YangFei1990 approved these changes Aug 12, 2026

View reviewed changes

Merge branch 'main' into ep_mxfp8

f1c48d5

phu0ngng merged commit 56a05e7 into NVIDIA:main Aug 12, 2026
10 of 14 checks passed

phu0ngng deleted the ep_mxfp8 branch August 12, 2026 02:05

fheinecke mentioned this pull request Aug 21, 2026

Fix PyTorch EP tests sharing a single timeout #3413

Open

13 tasks

Uh oh!

Conversation

phu0ngng commented Jul 28, 2026

Description

Type of change

Changes

PyTorch frontend (transformer_engine/pytorch/ep.py, distributed.py, csrc/extensions/ep.cpp)**

Common backend (common/ep/ep_backend.cpp, include/.../ep.h, comm_window.h)**

NCCL EP submodule**

Tests (tests/cpp_distributed/test_ep.cu, tests/pytorch/distributed/run_ep.py, run_test_ep.sh)**

Checklist:

Uh oh!

greptile-apps Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

phu0ngng commented Jul 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

phu0ngng commented Aug 6, 2026

Uh oh!

phu0ngng commented Aug 6, 2026

Uh oh!

phu0ngng commented Aug 7, 2026

Uh oh!

YangFei1990 left a comment

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented Aug 10, 2026

Uh oh!

phu0ngng commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

PyTorch frontend (`transformer_engine/pytorch/ep.py`, `distributed.py`, `csrc/extensions/ep.cpp`)**

Common backend (`common/ep/ep_backend.cpp`, `include/.../ep.h`, `comm_window.h`)**

Tests (`tests/cpp_distributed/test_ep.cu`, `tests/pytorch/distributed/run_ep.py`, `run_test_ep.sh`)**

greptile-apps Bot commented Jul 28, 2026 •

edited

Loading